import sys #import sys module
import numpy as np

$versionId: read_better.py 89 2013-05-03 11:35:17Z wilbertsmit@gmail.com $

#--------------------------------------------------

def is_line_ok( line,atom ):
    if line.startswith('ATOM') and line.split()[2] == atom:
        return True
    else:
        return False

#--------------------------------------------------

def grab_text( filename ):
    return open(filename,'r').readlines()

#--------------------------------------------------

def find_mean_pos( txt_lines,atom ):
    x_pos = []
    y_pos = []
    z_pos = []
    for line in txt_lines:
        if is_line_ok( line,atom ):
            line = line.split()
            x_pos.append( float(line[4]) )
            y_pos.append( float(line[5]) )
            z_pos.append( float(line[6]) )

    return np.mean(x_pos),np.mean(y_pos),np.mean(z_pos)

#--------------------------------------------------
def main( filename, atom_name ):
    txt_lines = grab_text( filename )
    ave_x,ave_y,ave_z = find_mean_pos( txt_lines, atom_name )

    print 'The mean position for ', atom_name, ' in, filename is:'
    print ave_x,ave_y,ave_z

#--------------------------------------------------
if __name__ == "__main__": #if statement
    if not len(sys.argv) == 3:
        print 'You must supply both a filename and a molecule name'
    else: #else staement
        main( sys.argv[1],sys.argv[2] )
